from PIL import Image
import pytesseract
from PIL import ImageFilter
from datetime import datetime
from core.capture_screen import capture_screen
import os
# 从图图片中取到 问题和答案
imgPathname = 'text_area.png'
def get_text_from_image(crop_area, directory = '.', compress_level = 1):
    print("capture time: ", datetime.now().strftime("%H:%M:%S"))
    screenshot_filename = "screenshot.png"
    save_text_area = os.path.join('.', imgPathname)
    capture_screen(screenshot_filename, directory)
    imgList = cropQAImg(os.path.join(directory, screenshot_filename), compress_level, crop_area)
    return imgList,ocr_img(imgList)
def cropQAImg(source_file, compress_level, crop_area):
    """
    crop the answer area

    :return:
    """
    image = Image.open(source_file)
    width, height = image.size[0], image.size[1]
    imageList = []
    for i in range(len(crop_area)):
        if i % 4 == 0:
            region = image.crop((width * crop_area[i], height * crop_area[i+1], width * crop_area[i+2], height * crop_area[i+3]))
            if compress_level == 1:
                region = region.convert("L")
            elif compress_level == 2:
                region = region.convert("1")
            imageList.append(region)

    print("screen 开始截成四个区域")
    return imageList
# 二值化算法
def binarizing(img,threshold):
    pixdata = img.load()
    w, h = img.size
    for y in range(h):
        for x in range(w):
            if pixdata[x, y] < threshold:
                pixdata[x, y] = 0
            else:
                pixdata[x, y] = 255
    return img


# 去除干扰线算法
def depoint(img):
    pixdata = img.load()
    w,h = img.size
    for y in range(1,h-1):
        for x in range(1,w-1):
            count = 0
            if pixdata[x,y-1] > 200:
                count = count + 1
            if pixdata[x,y+1] > 200:
                count = count + 1
            if pixdata[x-1,y] > 200:
                count = count + 1
            if pixdata[x+1,y] > 250:
                count = count + 1
            if count > 2:
                pixdata[x,y] = 255
    return img

def img_to_text (image):
    # 边缘增强滤波,不一定适用
    img = image.filter(ImageFilter.EDGE_ENHANCE)
    # 把图片变成二值图像
    question_im = binarizing(img, 190)
    question_im = depoint(question_im)
    # win环境
    # tesseract 路径
    # pytesseract.pytesseract.tesseract_cmd = 'C:\\Program Files (x86)\\Tesseract-OCR\\tesseract'
    # 语言包目录和参数
    # tessdata_dir_config = '--tessdata-dir "C:\\Program Files (x86)\\Tesseract-OCR\\tessdata" --psm 6'

    # mac 环境 记得自己安装训练文件
    # tesseract 路径
    # pytesseract.pytesseract.tesseract_cmd = '/usr/local/Cellar/tesseract/3.05.01/bin/tesseract'
    # 语言包目录和参数
    # tessdata_dir_config = '--tessdata-dir "/usr/local/Cellar/tesseract/3.05.01/share/tessdata/" --psm 6'

    # ubuntu 16.04 环境 记得自己安装训练文件 找不到位置就 whereis tesseract
    # tesseract 路径
    pytesseract.pytesseract.tesseract_cmd = '/usr/local/bin/tesseract'
    # 语言包目录和参数
    tessdata_dir_config = '--tessdata-dir "/usr/local/share/tessdata/" --psm 6'

    # lang 指定中文简体
    # question = pytesseract.image_to_string(question_im, lang='chi_sim', config=tessdata_dir_config)
    # lang 指定中文简体 不使用训练
    text = pytesseract.image_to_string(question_im, lang='chi_sim', config=tessdata_dir_config)
    text = text.replace("\n", "")[2:]
    # 处理将"一"识别为"_"的问题
    text = text.replace("_", "一")
    print('ocr:', text)
    return text

def ocr_img(imgList = []):
    textList = []
    for i, keyword in enumerate(imgList):
        text = img_to_text(keyword).strip()
        if i == 0 and text.endswith('?') == False:
            text = text + '?'
        textList.append(text)
    return textList

if __name__ == '__main__':
    image = Image.open("./test.png")
    image = [image]
    question,choices = get_text_from_image(image)

    print("识别结果:")
    print(question)
    print(choices)